#define PY_SSIZE_T_CLEAN
#include <Python.h>

#include <stdio.h>
#include <fcntl.h>

#include "wbfs-interface.h"
#include "lib-std.h"
#include "dclib/dclib-types.h"
#include "version.h"


/* Globals from libwwt (ones we use directly)
 * ---------------------------------
 *
 * wbfs-interface.h
 * ================
 * 
 * extern PartitionInfo_t *  first_partition_info
 *
 * iso-interface.h
 * ===============
 *
 * extern wd_select_t part_selector;
 * */


static PyObject * WwtError;


// -- Init WWT
// Straight from project/src/std-lib.c:SetupLib
// We don't call SetupLib() directly because it also does various things we
// dont need.
static void init_wwt() {
  //----- setup common key
  {
    static u8 key_base[] = "Wiimms WBFS Tool";
    static u8 key_mask[WD_CKEY__N][WII_KEY_SIZE] = {
        {0x25, 0x00, 0x94, 0x12, 0x2a, 0x3e, 0x4e, 0x1f, 0x03, 0x2d, 0xc5, 0xc9,
         0x30, 0x2a, 0x58, 0xab},
        {0xad, 0x5c, 0x95, 0x84, 0x80, 0xda, 0x93, 0xd5, 0x58, 0x06, 0xfe, 0x77,
         0xf9, 0xe7, 0x69, 0x22},
#ifdef SUPPORT_CKEY_DEVELOP
        {0x6f, 0x84, 0xf4, 0x5a, 0x05, 0x98, 0x68, 0xd2, 0xe5, 0x7f, 0xec, 0xbe,
         0x8b, 0xbd, 0x0e, 0xf6},
#endif
    };

    u8 h1[WII_HASH_SIZE], h2[WII_HASH_SIZE], key[WII_KEY_SIZE];

    SHA1(key_base, WII_KEY_SIZE, h1);
    int ci;
    for (ci = 0; ci < WD_CKEY__N; ci++) {
      SHA1(wd_get_common_key(ci), WII_KEY_SIZE, h2);

      int i;
      for (i = 0; i < WII_KEY_SIZE; i++) {
        key[i] = key_mask[ci][i] ^ h1[i] ^ h2[i];
      }

      wd_set_common_key(ci, key);
    }
  }

  //----- setup data structures
  wd_initialize_select(&part_selector);
}


// -- Raise exception from enumError
static void raise_wit_exception(enumError err, const char * methodName) {
  const char * errName = GetErrorName(err, "Unknown error name");
  const char * errText = GetErrorText(err, "Unknown error text");
  PyErr_Format(WwtError, "%s returned a error: %s - %s",
               methodName, errName, errText);
}


// -- Internal function to open a wbfs partition.
static int open_wbfs_partition(const char * partition, WBFS_t * wbfs, bool open_modify) {
  // NOTE: Always make sure ResetWBFS is called after this returned 0
  enumError err;
  PartitionInfo_t * info;

  // Create partition info for the requested partition/file
  // @NOTE: CreatePartitionInfo() will override existing ones for us and create
  //        new partition info for new ones.
  CreatePartitionInfo(partition, PS_PARAM);

  // Search for the partition info
  info = first_partition_info;
  while (info && strcmp(info->real_path, partition)) {
    info = info->next;
  }

  // Partition info not found.
  if (info == NULL) {
    PyErr_SetString(WwtError, "Failed to create partition info");
    return 1;
  }

  // Update partition info
  // @NOTE: AnalyzePartitions() will for example set the 'part_mode'.
  err = AnalyzePartitions(stderr, true, true);
  if (err) {
    raise_wit_exception(err, "AnalyzePartitions");
    return 1;
  }

  // Check if its a WBFS partition we are dealing with.
  if (info->part_mode != PM_WBFS) {
    PyErr_SetString(WwtError, "Not a valid WBFS partition");
    return 1;
  }

  InitializeWBFS(wbfs);
  err = OpenPartWBFS(wbfs, info, open_modify);
  if (err) {
    raise_wit_exception(err, "OpenPartWBFS");
    return 1;
  }

  return 0;
}


static int find_slot_by_id6(WBFS_t * wbfs, ccp id6) {
  int slot;
  int foundSlot = -1;
  ccp slotId6;
  id6_t * idList;

  idList = wbfs_load_id_list(wbfs->wbfs, false);

  for (slot = 0; slot < wbfs->wbfs->max_disc; slot++) {
    // Get the ID6 for this slot
    slotId6 = idList[slot];

    // empty or invalid slot?
    if (!*slotId6) {
      continue;
    }

    if (strncmp(id6, slotId6, 6) == 0) {
      foundSlot = slot;
      break;
    }
  }

  return foundSlot;
}


// -- Get version
static PyObject * method_version(PyObject *self, PyObject *args) {
  return Py_BuildValue("{s:s, s:s, s:s}",
                       "version", VERSION,
                       "rev", REVISION,
                       "arch", SYSTEM2);
}


// -- Returns a list with WBFS formatted device patitions
static PyObject * method_find_wbfs_partitions(PyObject * self, PyObject * args) {
  PyObject * list;

  // Reset WIT internals
  PartitionInfo_t * info;
  for (info=first_partition_info; info; info=info->next) {
    info->part_mode = PM_UNKNOWN; // this will trigger a re-read on
                                  // AnalyzePartitions so wbfs partitions will
                                  // be found after format from non wbfs
                                  // partition.
  }

  ScanPartitions(true);
  enumError err = AnalyzePartitions(stderr,true,true);

  if (err) {
    raise_wit_exception(err, "AnalyzePartitions");
    return NULL;
  }

  list  = PyList_New(0);

  for (info = first_partition_info; info; info = info->next) {
    if (info->part_mode == PM_WBFS) {
      PyObject * pathStr = Py_BuildValue("s", info->path);
      PyList_Append(list, pathStr); // TODO check return value
    }
  }
  return list;
}


// -- Get information about a WBFS partition and its contents.
static PyObject * method_get_wbfs_partition_info(PyObject * self, PyObject * args, PyObject * kwargs) {
  WBFS_t wbfs;
  const char * partition = NULL;
  int disc_count, max_discs, used_mib, free_mib, wbfs_version;
  PyObject * discList = NULL;

  // parse args/kwargs
  bool includeDiscs = false;
  static char * kwlist[] = {"partition", "discs", NULL};
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s|b", kwlist, &partition, &includeDiscs)) {
    PyErr_SetString(WwtError, "Invalid arguments");
    return NULL;
  }

  // open partition
  if (open_wbfs_partition(partition, &wbfs, false) != 0) {
    return NULL;
  }


  WDiscList_t * wlist = GenerateWDiscList(&wbfs, 0);
  if (wlist) {

    discList = PyList_New(0);
    if (includeDiscs == true) {
      SortWDiscList(wlist, sort_mode, SORT_TITLE, false);

      // Go over all discs
      WDiscListItem_t * witem;
      int i;
      for (i = wlist->used, witem = wlist->first_disc; i-- > 0; witem++ ) {
        /* {
         *   "id6"   : "",
         *   "title" : "",
         *   "region": "",
         *   "slot"  : 0,
         *   "size" : 0,
         *   "ftype" : 0
         * }*/
        PyObject * discObj = Py_BuildValue("{s:s, s:s, s:i, s:i, s:i, s:i}",
                                           "id6"   , witem->id6,
                                           "title" , (witem->title != NULL) ? witem->title : witem->name64,
                                           "region", GetRegionInfo(witem->id6[3])->reg,
                                           "slot"  , (int)witem->wbfs_slot,
                                           "size"  , (int)witem->size_mib,
                                           "ftype" , (int)witem->ftype);

        PyList_Append(discList, discObj); // TODO check return value
      }
    }

    disc_count   = wlist->used;
    max_discs    = wbfs.total_discs;
    used_mib     = wlist->total_size_mib;
    free_mib     = wbfs.free_mib;
    wbfs_version = (int)wbfs.wbfs->head->wbfs_version;

    FreeWDiscList(wlist);
  }

  // close partition
  ResetWBFS(&wbfs);

  if (!wlist) {
    PyErr_SetString(WwtError, "Could not generate disc list");
    return NULL;
  }

  return Py_BuildValue("{s:i, s:i, s:i, s:i, s:i, s:O}",
                       "discs"       , disc_count,
                       "discs_max"   , max_discs,
                       "used"        , used_mib,
                       "free"        , free_mib,
                       "wbfs_version", wbfs_version,
                       "discs"       , discList);
}


// -- Check a WBFS partition on errors
static PyObject * method_check_wbfs_partition(PyObject * self, PyObject * args) {
  int errPy;
  WBFS_t wbfs;
  CheckWBFS_t ck;
  const char * partition = NULL;
  PyObject * resultObj = NULL;
  PyObject * errorsObj = NULL;

  // parse the 'partition' argument
  errPy = PyArg_ParseTuple(args, "s", &partition);
  if (!errPy) {
    PyErr_SetString(WwtError, "No partition argument");
    return NULL;
  }

  // Reset WIT internals else the read-out may be invalid when the partition
  // has been changed by a another instance or program.
  CloseWBFSCache();

  // open wbfs partition
  if (open_wbfs_partition(partition, &wbfs, false) != 0) {
    return NULL;
  }

  /*{
    "status"            : 0,  # CheckWBFS_t.err                  OK | WARNING | WBFS_INVALID
    "error_count"       : 0,  # CheckWBFS_t.err_total
    "invalid_disc_count": 0,  # CheckWBFS_t.invalid_disc_count
    "errors": {
      "used_blocks": 0,       # CheckWBFS_t.err_fbt_used         used WBFS sectors marked as free
      "free_blocks": 0,       # CheckWBFS_t.err_fbt_free         free WBFS sectors marked as used
      "free_wbfs0" : 0,       # CheckWBFS_t.err_fbt_free_wbfs0   based on a bug in libwbfs v0
      "overlap"    : 0,       # CheckWBFS_t.err_bl_overlap       WBFS sectors are used by 2 or more discs
      "references" : 0,       # CheckWBFS_t.err_bl_invalid       invalid WBFS block references
      "disc_blocks": 0        # CheckWBFS_t.err_no_blocks        discs with no valid WBFS blocks
    }
  }*/

  InitializeCheckWBFS(&ck);
  // NOTE: FIXME Assumption: when err >= ERR_ERROR there is a internal error
  if (CheckWBFS(&ck, &wbfs, 0, 0, NULL, 0) < ERR_ERROR) {
    errorsObj = Py_BuildValue("{s:i, s:i, s:i, s:i, s:i, s:i}",
      "used_blocks", ck.err_fbt_used,
      "free_blocks", ck.err_fbt_free,
      "free_wbfs0" , ck.err_fbt_free_wbfs0,
      "overlap"    , ck.err_bl_overlap,
      "ref_blocks" , ck.err_bl_invalid,
      "disc_blocks", ck.err_no_blocks);

    resultObj = Py_BuildValue("{s:i, s:i, s:i, s:O}",
      "status"            , ck.err,
      "error_count"       , ck.err_total,
      "invalid_disc_count", ck.invalid_disc_count,
      "errors"            , errorsObj);
  }
  ResetCheckWBFS(&ck);

  // close partition
  ResetWBFS(&wbfs);

  if (resultObj != NULL) {
    return resultObj;
  }

  PyErr_SetString(WwtError, "CheckWBFS failed");
  return NULL;
}


// -- Attempt to repair a WBFS partition
static PyObject * method_repair_wbfs_partition(PyObject * self, PyObject * args, PyObject * kwargs)
{
  enumError     err;
  WBFS_t        wbfs;
  CheckWBFS_t   ck;
  const char  * partition = NULL;
  int repairMode = REPAIR_ALL; // default
  bool testMode  = true;       // default

  // parse args/kwargs
  static char * kwlist[] = {"partition", "repairMode", "testMode", NULL};
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s|ib", kwlist, &partition, &repairMode, &testMode)) {
    PyErr_SetString(WwtError, "Invalid arguments");
    return NULL;
  }

  // open wbfs partition in writing mode
  if (open_wbfs_partition(partition, &wbfs, true) != 0)
  {
    return NULL;
  }

  err = 0;

  InitializeCheckWBFS(&ck);
  // NOTE: FIXME Assumption: when err >= ERR_ERROR there is a internal error
  if (CheckWBFS(&ck, &wbfs, 0, 0, NULL, 0) < ERR_ERROR
      || (repairMode & REPAIR_FBT && wbfs.wbfs->head->wbfs_version < WBFS_VERSION)
      || (repairMode & REPAIR_INODES && ck.no_iinfo_count))
  {
    if (ck.err)
    {
      err = RepairWBFS(&ck, (int)testmode, repairMode, -1, NULL, 0);
    }
  }
  ResetCheckWBFS(&ck);

  // close partition
  ResetWBFS(&wbfs);

  if (!err) {
    Py_RETURN_TRUE;
  }

  // repairing failed
  Py_RETURN_FALSE;
}


// -- Extract a disc from a WBFS partition.
static PyObject * method_extract_from_wbfs(PyObject * self, PyObject * args, PyObject * kwargs)
{
  enumError err;
  WBFS_t wbfs;
  SuperFile_t fo;      // file out
  wd_header_t * dhead; // disc header
  wd_disc_t * disc = NULL;

  // args/kwargs
  const char * partition = NULL;
  const char * id6       = NULL;
  const char * dest      = NULL;
  bool testMode          = true; // default

  // parse args/kwargs
  static char * kwlist[] = {"partition", "id6", "dest", "testMode", NULL};
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "sss|b", kwlist, &partition, &id6, &dest, &testMode)) {
    PyErr_SetString(WwtError, "Invalid arguments");
    return NULL;
  }

  // open wbfs partition read only
  if (open_wbfs_partition(partition, &wbfs, false) != 0)
  {
    return NULL;
  }

  // Game not found
  if (find_slot_by_id6(&wbfs, id6) == -1) {
    ResetWBFS(&wbfs);
    PyErr_SetString(WwtError, "No slot for given ID6 found");
    return NULL;
  }

  // open the disc on the wbfs partition
  err = OpenWDiscID6(&wbfs, id6);

  // failed to open the disc on the wbfs partition
  if (err) {
    ResetWBFS(&wbfs);
    raise_wit_exception(err, "OpenWDiscID6");
    return NULL;
  }

  // get the disc header
  dhead = GetWDiscHeader(&wbfs);

  // failed the get disc header
  if (!dhead) {
    CloseWDisc(&wbfs);
    ResetWBFS(&wbfs);
    PyErr_SetString(WwtError, "Failed the get disc header for given ID6");
    return NULL;
  }

  if (testMode) {
    printf("Would have extracted %s '%s' to '%s'\n", id6, dhead->disc_title, dest);
    ResetWBFS(&wbfs);
    CloseWDisc(&wbfs);
    Py_RETURN_TRUE;
  }

  // TODO remove later
  printf("Going to extract %s '%s' to '%s'\n", id6, dhead->disc_title, dest);

  err = OpenWDiscSF(&wbfs);
  if (err) {
    ResetWBFS(&wbfs);
    raise_wit_exception(err, "OpenWDiscSF");
    return NULL;
  }

  disc = OpenDiscSF(wbfs.sf, true, true);
  if (disc == NULL) {
    ResetWBFS(&wbfs);
    PyErr_SetString(WwtError, "OpenDiscSF failed");
    return NULL;
  }

  Py_BEGIN_ALLOW_THREADS

  // Output file type
  enumOFT oft = CalcOFT(output_file_type, dest, 0, OFT__DEFAULT);

  if (oft == OFT_FST) {
    err = ExtractImage(wbfs.sf, dest, false, false);
  }
  else {
    InitializeSF(&fo);
    SetWFileName(&fo.f, dest, true);
    fo.f.create_directory = true;
    fo.src                = wbfs.sf;
    err = CopyImage(wbfs.sf, &fo, oft, true, false, false);
  }

  // close the disc and wbfs
  CloseWDisc(&wbfs);
  ResetWBFS(&wbfs);

  Py_END_ALLOW_THREADS

  if (err) {
    Py_RETURN_FALSE;
  }

  Py_RETURN_TRUE;
}


// -- Add a disc to a WBFS partition
static PyObject * method_add_to_wbfs(PyObject * self, PyObject * args, PyObject * kwargs)
{
  enumError err;
  WBFS_t wbfs;
  SuperFile_t inputFile;   // File descriptor for input gamefile.
  wd_disc_t * disc;        // Disc to add (from inputFile).
  WDiscList_t * wlist;     // For iterating over existing discs on the WBFS
                           // partition to check wheter the game is already
                           // present or not.
  WDiscListItem_t * witem; // Disc item of partition, see wlist.
  int i;                   // For disc iteration.

  // args/kwargs
  const char * partition = NULL;
  const char * gamefile  = NULL;
  bool testMode          = true;

  // parse args/kwargs
  static char * kwlist[] = {"partition", "gamefile", "testMode", NULL};
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "ss|b", kwlist, &partition, &gamefile, &testMode)) {
    PyErr_SetString(WwtError, "Invalid arguments");
    return NULL;
  }

  // open wbfs partition
  if (open_wbfs_partition(partition, &wbfs, true) != 0)
  {
    return NULL;
  }

  // open input file
  InitializeSF(&inputFile);
  //                                    _______ bool allow_non_iso
  //                                   /     __ bool open_modify
  //                                  /     /
  err = OpenSF(&inputFile, gamefile, true, false);

  // failed to open input file
  if (err) {
    ResetWBFS(&wbfs);
    raise_wit_exception(err, "OpenSF");
    return NULL;
  }

  disc = OpenDiscSF(&inputFile, true, true);
  if (disc == NULL) {
    ResetWBFS(&wbfs);
    CloseSF(&inputFile, 0);
    PyErr_SetString(WwtError, "OpenDiscSF failed");
    return NULL;
  }

  // TODO remove this
  printf("Input: %s %s\n", disc->dhead.id6.id6, disc->dhead.disc_title);

  // game already present
  if (find_slot_by_id6(&wbfs, disc->dhead.id6.id6) != -1) {
    ResetWBFS(&wbfs);
    CloseSF(&inputFile, 0);
    PyErr_SetString(WwtError, "Game already present");
    return NULL;
  }

  // testMode
  if (testMode) {
    ResetWBFS(&wbfs);
    CloseSF(&inputFile, 0);
    Py_RETURN_TRUE;
  }

  // add the disc
  Py_BEGIN_ALLOW_THREADS
  err = AddWDisc(&wbfs, &inputFile, &part_selector);
  Py_END_ALLOW_THREADS

  // cleanup
  CloseSF(&inputFile, 0);
  ResetWBFS(&wbfs);

  // adding the disc failed
  if (err) {
    Py_RETURN_FALSE;
  }

  // adding the disc succeeded
  Py_RETURN_TRUE;
}


// -- Remove a disc from a WBFS partition
static PyObject * method_remove_from_wbfs(PyObject * self, PyObject * args, PyObject * kwargs)
{
  enumError err;
  WBFS_t wbfs;
  int slot, foundSlot;

  // args/kwargs
  const char * partition = NULL;
  const char * id6       = NULL;
  bool testMode          = true;

  // parse args/kwargs
  static char * kwlist[] = {"partition", "id6", "testMode", NULL};
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "ss|b", kwlist, &partition, &id6, &testMode)) {
    PyErr_SetString(WwtError, "Invalid arguments");
    return NULL;
  }

  // open wbfs partition
  if (open_wbfs_partition(partition, &wbfs, true) != 0)
  {
    return NULL;
  }

  // find the slot for the given id6
  foundSlot = -1;
  for (slot = 0; slot < wbfs.wbfs->max_disc; slot++) {
    // Get the ID6 for this slot
    ccp slotId6 = wbfs_load_id_list(wbfs.wbfs, false)[slot];

    // empty or invalid slot?
    if (!*id6) {
      continue;
    }

    if (strncmp(id6, slotId6, 6) == 0) {
      foundSlot = slot;
      break;
    }
  }

  // slot not found
  if (foundSlot == -1) {
    ResetWBFS(&wbfs);
    PyErr_SetString(WwtError, "No slot for given ID6 found");
    return NULL;
  }

  if (testMode) {
    Py_RETURN_TRUE;
  }

  // remove the disc
  err = RemoveWDisc(&wbfs, id6, foundSlot, false);

  // cleanup
  ResetWBFS(&wbfs);

  // remove disc failed
  if (err) {
    Py_RETURN_FALSE;
  }

  Py_RETURN_TRUE;
}


// -- Create a new WBFS partition file
static PyObject * method_create_new_wbfs_file(PyObject * self, PyObject * args, PyObject * kwargs)
{
  // GiB from dclib/dclib-types.h
  enumError err;
  WFile_t f;
  off_t size;

  // args/kwargs
  const char * filepath;
  u32 inputSize;         // In GB, so give 1 for 1 GB
  u32 hss                = HD_SECTOR_SIZE; // HD Sector Size
  u32 wss                = 0; // WBFS Sector Size (should be set automatically
                              // when left to 0)
  bool testMode          = true;

  // parse args/kwargs
  static char * kwlist[] = {"filepath", "size", "hss", "wss", "testMode", NULL};
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "sI|IIb", kwlist, &filepath,
                                   &inputSize, &hss, &wss, &testMode))
  {
    PyErr_SetString(WwtError, "Invalid arguments");
    return NULL;
  }

  // Reset WIT internals (else when we create a new file, then delete it and
  // then create a new file again with the same name it will fail because the
  // wit internals keeps track of stuff.
  ResetStringField(&created_files);
  ResetStringField(&source_list);
  ResetStringField(&recurse_list);
  // Also invalidate WIT it's cache, else the second time we create a file with
  // the same filepath (and filename) the partition will be invalid.
  CloseWBFSCache();

  // set actual size in bytes
  size = (off_t)(inputSize * GiB);

  // Stat given filepath
  struct stat st;
  err = StatFile(&st, filepath, 0);

  // File does not exist (could not stat it, as we would aspect)
  if (err) {
    if (testMode) {
      printf("Would have created file '%s' [%ld]\n", filepath, size);
    }
    else {
      // Create new file
      InitializeWFile(&f);
      err = CreateWFile(&f, filepath, IOM_NO_STREAM, 1);

      // Create new file failed
      if (err) {
        raise_wit_exception(err, "CreateWFile");
        return NULL;
      }

      // Set file size of the new file
      err = SetSizeF(&f, size);

      // Setting file size of the new file failed
      if (err) {
        ClearWFile(&f, true);
        raise_wit_exception(err, "SetSizeF");
        return NULL;
      }

      // cleanup
      ClearWFile(&f, false);
    }
  }
  // File already exists
  else {
    PyErr_SetString(WwtError, "File already exists!");
    return NULL;
  }

  if (testMode) {
    printf("Would have formatted file '%s' [%ld]\n", filepath, size);
    Py_RETURN_TRUE;
  }

  // Prepare format of new file
  WBFS_t wbfs;
  InitializeWBFS(&wbfs);

  wbfs_param_t par;
  memset(&par, 0, sizeof(par));
  par.hd_sector_size   = hss;
  par.wbfs_sector_size = wss;
  par.reset            = 1;
  par.clear_inodes     = true;
  par.setup_iinfo      = false;

  // Format new file
  err = FormatWBFS(&wbfs, filepath, true, &par, 0, 0);

  // Clean up
  ResetWBFS(&wbfs);

  // Format of new file failed
  if (err) {
    printf("Format of '%s' failed!\n", filepath);
    Py_RETURN_FALSE;
  }

  // Success!
  Py_RETURN_TRUE;
}


// -- (Re-)Format a partition to WBFS
static PyObject * method_format_to_wbfs(PyObject * self, PyObject * args, PyObject * kwargs)
{
  enumError err;
  struct stat st;
  WFile_t f;
  WBFS_t wbfs;
  AWData_t awd;
  wbfs_param_t par;
  u32 dev_hss;

  // args/kwargs
  ccp partition;
  u32 hss                = 0; // HD Sector Size
  u32 wss                = 0; // WBFS Sector Size (should be set automatically
                              // when left to 0)
  bool recover           = false;
  bool testMode          = true;

  // parse args/kwargs
  static char * kwlist[] = {"partition", "hss", "wss", "recover", "testMode", NULL};
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s|IIbb", kwlist, &partition,
                                   &hss, &wss, &recover, &testMode))
  {
    PyErr_SetString(WwtError, "Invalid arguments");
    return NULL;
  }

  // stat partition
  err = StatFile(&st, partition, 0);

  if (err) {
    raise_wit_exception(err, "StatFile");
    return NULL;
  }

  // Set HSS and WSS from AnalyzeWBFS (When HSS or WSS is not given)
  if (recover && (!hss || !wss)) {
    InitializeWFile(&f); // TODO move out this branch?
    err = OpenWFile(&f, partition, IOM_IS_WBFS_PART);
    if (err == ERR_OK) {
      AnalyzeWBFS(&awd, &f);
      if (awd.n_record) {
        hss = awd.rec[0].hd_sector_size;
        wss = awd.rec[0].wbfs_sector_size;
      }
      ResetWFile(&f, 0);
    }
  }

  // Set HD Sector Size when not set for block or char devices.
  if (!hss && (S_ISBLK(st.st_mode) || S_ISCHR(st.st_mode))) {
    const int fd = open(partition, O_RDONLY);
    if (fd != -1) {
      dev_hss = GetHSS(fd, 0);
      close(fd);
      if (dev_hss) {
        hss = HD_SECTOR_SIZE;
        while (hss && hss < dev_hss) {
          hss <<= 1;
        }
      }
    }
  }

  // Set default HSS when HSS is smaller then the default
  if (hss < HD_SECTOR_SIZE) {
    hss = HD_SECTOR_SIZE;
  }

  // TODO WSS should be 2 times HSS?
  if (wss && wss <= hss) {
    wss = 2 * hss;
  }

  if (testMode) {
    Py_RETURN_TRUE;
  }

  // Prepare format
  InitializeWBFS(&wbfs);

  memset(&par, 0, sizeof(par));
  par.hd_sector_size   = hss;
  par.wbfs_sector_size = wss;
  par.reset            = 1;
  par.clear_inodes     = !recover;
  par.setup_iinfo      = !recover;

  // Actual format
  err = FormatWBFS(&wbfs, partition, true, &par, 0, 0);

  if (err == ERR_OK) {
    // Recover
    if (recover) {
      err = RecoverWBFS(&wbfs, partition, false);
    }

    // Clean up
    ResetWBFS(&wbfs);

    // Recover failed
    if (err) {
      raise_wit_exception(err, "RecoverWBFS");
      return NULL;
    }

    Py_RETURN_TRUE;
  }

  raise_wit_exception(err, "FormatWBFS");
  return NULL;
}


// -- Get information on a file
static PyObject * method_get_file_info(PyObject * self, PyObject * args)
{
  enumError err;
  const char * filePath;
  SuperFile_t inputFile;
  wd_disc_t * disc;
  bool isSupported;
  enumFileType fType;
  ccp id6;
  ccp title;
  int region;

  // parse the 'filePath' argument
  if (!PyArg_ParseTuple(args, "s", &filePath)) {
    PyErr_SetString(WwtError, "No filePath argument");
    return NULL;
  }

  // open input file
  InitializeSF(&inputFile);
  //                                    _______ bool allow_non_iso
  //                                   /     __ bool open_modify
  //                                  /     /
  err = OpenSF(&inputFile, filePath, true, false);

  // failed to open input file
  if (err) {
    raise_wit_exception(err, "OpenSF");
    return NULL;
  }

  disc = OpenDiscSF(&inputFile, true, true);
  if (disc == NULL) {
    CloseSF(&inputFile, 0);
    return Py_BuildValue("{s:i, s:i, s:s, s:s, s:i}",
                         "supported", false,
                         "type"     , 0,
                         "id6"      , "",
                         "title"    , "",
                         "region"   , 0);
  }

  isSupported = false;
  fType       = inputFile.f.ftype;
  id6         = disc->dhead.id6.id6;
  title       = disc->dhead.disc_title;
  region      = GetRegionInfo(id6[3])->reg;
  if (fType && FT_A_WDISC || fType && FT_A_WII_ISO) {
    isSupported = true;
  }

  PyObject * result = Py_BuildValue("{s:i, s:i, s:s, s:s, s:i}",
                                    "supported", isSupported,
                                    "type"     , fType,
                                    "id6"      , id6,
                                    "title"    , title,
                                    "region"   , region);

  CloseSF(&inputFile, 0);

  return result;
}


// -- Copy a game from a WBFS partition to another WBFS partition
static PyObject * method_wbfs_to_wbfs(PyObject * self, PyObject * args)
{
  enumError err;
  WBFS_t fromWbfs;
  WBFS_t toWbfs;
  wd_disc_t * disc;
  WDiscList_t * wlist;
  WDiscListItem_t * witem;
  int i;

  // args
  ccp fromPart;
  ccp toPart;
  ccp id6;

  // parse the 'filePath' argument
  if (!PyArg_ParseTuple(args, "sss", &fromPart, &toPart, &id6)) {
    PyErr_SetString(WwtError, "Invalid arguments");
    return NULL;
  }

  if (strlen(id6) > 6) {
    PyErr_SetString(WwtError, "Given ID6 is to long!");
    return NULL;
  }

  /////////////////////////////////////////
  // Open/prepare the source WBFS partition
  /////////////////////////////////////////

  // open wbfs fromPart read only
  if (open_wbfs_partition(fromPart, &fromWbfs, false) != 0)
  {
    return NULL;
  }

  // Game not found
  if (find_slot_by_id6(&fromWbfs, id6) == -1) {
    ResetWBFS(&fromWbfs);
    PyErr_SetString(WwtError, "No slot for given ID6 found");
    return NULL;
  }

  // open the disc on the wbfs partition
  err = OpenWDiscID6(&fromWbfs, id6);

  // failed to open the disc on the wbfs partition
  if (err) {
    ResetWBFS(&fromWbfs);
    raise_wit_exception(err, "OpenWDiscID6");
    return NULL;
  }

  err = OpenWDiscSF(&fromWbfs);
  if (err) {
    CloseWDisc(&fromWbfs);
    ResetWBFS(&fromWbfs);
    raise_wit_exception(err, "OpenWDiscSF");
    return NULL;
  }

  disc = OpenDiscSF(fromWbfs.sf, true, true);
  if (disc == NULL) {
    CloseWDisc(&fromWbfs);
    ResetWBFS(&fromWbfs);
    PyErr_SetString(WwtError, "OpenDiscSF");
    return NULL;
  }

  /////////////////////////////////////////
  // Open/prepare the target WBFS partition
  /////////////////////////////////////////

  // open wbfs toPart write mode
  if (open_wbfs_partition(toPart, &toWbfs, true) != 0)
  {
    ResetWBFS(&fromWbfs);
    return NULL;
  }

  // game already present
  if (find_slot_by_id6(&toWbfs, id6) != -1) {
    ResetWBFS(&fromWbfs);
    ResetWBFS(&toWbfs);
    PyErr_SetString(WwtError, "Game already present");
    return NULL;
  }

  ////////////////
  // Copy the game
  ////////////////

  Py_BEGIN_ALLOW_THREADS
  err = AddWDisc(&toWbfs, fromWbfs.sf, &part_selector);

  // cleanup
  ResetWBFS(&fromWbfs);
  ResetWBFS(&toWbfs);

  Py_END_ALLOW_THREADS

  // copy the game failed
  if (err) {
    Py_RETURN_FALSE;
  }

  // copying the game succeeded
  Py_RETURN_TRUE;
}


static PyMethodDef libwwtMethods[] = {
  // pywwt.version()
  {"version", method_version, METH_VARARGS,
    "Returns the version of wwt."},

  // pywwt.find_wbfs_partitions()
  {"find_wbfs_partitions", method_find_wbfs_partitions, METH_VARARGS,
    "Find WBFS partitions."},

  // pywwt.get_wbfs_partition_info()
  {"get_wbfs_partition_info", (PyCFunction)method_get_wbfs_partition_info, METH_VARARGS | METH_KEYWORDS,
    "Get WBFS partition info"},

  // pywwt.check_wbfs_partition()
  {"check_wbfs_partition", method_check_wbfs_partition, METH_VARARGS,
    "Check WBFS partition on errors"},

  // pywwt.repair_wbfs_partition()
  {"repair_wbfs_partition", (PyCFunction)method_repair_wbfs_partition, METH_VARARGS | METH_KEYWORDS,
    "Repair WBFS partition"},

  // pywwt.extract_from_wbfs()
  {"extract_from_wbfs", (PyCFunction)method_extract_from_wbfs, METH_VARARGS | METH_KEYWORDS,
    "Extract disc from WBFS"},

  // pywwt.add_to_wbfs()
  {"add_to_wbfs", (PyCFunction)method_add_to_wbfs, METH_VARARGS | METH_KEYWORDS,
    "Add disc to WBFS"},

  // pywwt.remove_from_wbfs()
  {"remove_from_wbfs", (PyCFunction)method_remove_from_wbfs, METH_VARARGS | METH_KEYWORDS,
    "Remove disc from WBFS"},

  // pywwt.create_new_wbfs_file()
  {"create_new_wbfs_file", (PyCFunction)method_create_new_wbfs_file, METH_VARARGS | METH_KEYWORDS,
    "Create a new WBFS partition file"},

  // pywwt.format_to_wbfs()
  {"format_to_wbfs", (PyCFunction)method_format_to_wbfs, METH_VARARGS | METH_KEYWORDS,
    "Format a partition to WBFS"},

  // pywwt.get_file_info()
  {"get_file_info", method_get_file_info, METH_VARARGS,
    "Get information on a file"},

  // pywwt.wbfs_to_wbfs()
  {"wbfs_to_wbfs", method_wbfs_to_wbfs, METH_VARARGS,
    "Copy a game from a WBFS partiton to another WBFS partition"},

  {NULL, NULL, 0, NULL}
};


static struct PyModuleDef libwwtmodule = {
  PyModuleDef_HEAD_INIT,
  "pywwt",
  "Brings the power of WWT (Wiimms WBFS Tool) to Python",
  -1,
  libwwtMethods
};


PyMODINIT_FUNC PyInit_pywwt(void) {
  PyObject * module = PyModule_Create(&libwwtmodule);

  // https://docs.python.org/3/c-api/exceptions.html#exception-classes
  WwtError = PyErr_NewException("pywwt.error", NULL, NULL);
  Py_XINCREF(WwtError);
  if (PyModule_AddObject(module, "error", WwtError) < 0) {
    Py_XDECREF(WwtError);
    Py_CLEAR(WwtError);
    Py_DECREF(module);
    return NULL;
  }

  // src/lib-std.h enum enumFileType
  PyModule_AddIntConstant(module, "FT_UNKNOWN", FT_UNKNOWN);
  PyModule_AddIntConstant(module, "FT_ID_DIR", FT_ID_DIR);
  PyModule_AddIntConstant(module, "FT_ID_FST", FT_ID_FST);
  PyModule_AddIntConstant(module, "FT_ID_WBFS", FT_ID_WBFS);
  PyModule_AddIntConstant(module, "FT_ID_GC_ISO", FT_ID_GC_ISO);
  PyModule_AddIntConstant(module, "FT_ID_WII_ISO", FT_ID_WII_ISO);
  PyModule_AddIntConstant(module, "FT_ID_DOL", FT_ID_DOL);
  PyModule_AddIntConstant(module, "FT_ID_SIG_BIN", FT_ID_SIG_BIN);
  PyModule_AddIntConstant(module, "FT_ID_CERT_BIN", FT_ID_CERT_BIN);
  PyModule_AddIntConstant(module, "FT_ID_TIK_BIN", FT_ID_TIK_BIN);
  PyModule_AddIntConstant(module, "FT_ID_TMD_BIN", FT_ID_TMD_BIN);
  PyModule_AddIntConstant(module, "FT_ID_HEAD_BIN", FT_ID_HEAD_BIN);
  PyModule_AddIntConstant(module, "FT_ID_BOOT_BIN", FT_ID_BOOT_BIN);
  PyModule_AddIntConstant(module, "FT_ID_FST_BIN", FT_ID_FST_BIN);
  PyModule_AddIntConstant(module, "FT_ID_PATCH", FT_ID_PATCH);
  PyModule_AddIntConstant(module, "FT_ID_OTHER", FT_ID_OTHER);
  PyModule_AddIntConstant(module, "FT_A_ISO", FT_A_ISO);
  PyModule_AddIntConstant(module, "FT_A_GC_ISO", FT_A_GC_ISO);
  PyModule_AddIntConstant(module, "FT_A_WII_ISO", FT_A_WII_ISO);
  PyModule_AddIntConstant(module, "FT_A_WDISC", FT_A_WDISC);
  PyModule_AddIntConstant(module, "FT_A_WDF1", FT_A_WDF1);
  PyModule_AddIntConstant(module, "FT_A_WDF2", FT_A_WDF2);
  PyModule_AddIntConstant(module, "FT_A_WIA", FT_A_WIA);
  PyModule_AddIntConstant(module, "FT_A_CISO", FT_A_CISO);
  PyModule_AddIntConstant(module, "FT_A_GCZ", FT_A_GCZ);
  PyModule_AddIntConstant(module, "FT_A_NKIT_ISO", FT_A_NKIT_ISO);
  PyModule_AddIntConstant(module, "FT_A_NKIT_GCZ", FT_A_NKIT_GCZ);
  PyModule_AddIntConstant(module, "FT_A_REGFILE", FT_A_REGFILE);
  PyModule_AddIntConstant(module, "FT_A_BLOCKDEV", FT_A_BLOCKDEV);
  PyModule_AddIntConstant(module, "FT_A_CHARDEV", FT_A_CHARDEV);
  PyModule_AddIntConstant(module, "FT_A_SEEKABLE", FT_A_SEEKABLE);
  PyModule_AddIntConstant(module, "FT_A_WRITING", FT_A_WRITING);
  PyModule_AddIntConstant(module, "FT_A_PART_DIR", FT_A_PART_DIR);
  PyModule_AddIntConstant(module, "FT_M_WDF", FT_M_WDF);
  PyModule_AddIntConstant(module, "FT_M_NKIT", FT_M_NKIT);

  // src/lib-std.h enum enumFileType
  PyModule_AddIntConstant(module, "REPAIR_NONE", REPAIR_NONE);
  PyModule_AddIntConstant(module, "REPAIR_FBT", REPAIR_FBT);
  PyModule_AddIntConstant(module, "REPAIR_INODES", REPAIR_INODES);
  PyModule_AddIntConstant(module, "REPAIR_DEFAULT", REPAIR_DEFAULT);
  PyModule_AddIntConstant(module, "REPAIR_RM_INVALID", REPAIR_RM_INVALID);
  PyModule_AddIntConstant(module, "REPAIR_RM_OVERLAP", REPAIR_RM_OVERLAP);
  PyModule_AddIntConstant(module, "REPAIR_RM_FREE", REPAIR_RM_FREE);
  PyModule_AddIntConstant(module, "REPAIR_RM_EMPTY", REPAIR_RM_EMPTY);
  PyModule_AddIntConstant(module, "REPAIR_RM_ALL", REPAIR_RM_ALL);
  PyModule_AddIntConstant(module, "REPAIR_ALL", REPAIR_ALL);

  // src/patch.h
  PyModule_AddIntConstant(module, "REGION_JAP", REGION_JAP);
  PyModule_AddIntConstant(module, "REGION_USA", REGION_USA);
  PyModule_AddIntConstant(module, "REGION_EUR", REGION_EUR);
  PyModule_AddIntConstant(module, "REGION_KOR", REGION_KOR);

  init_wwt();

  return module;
}
